#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
tts_client.py  (Piper-only)

特点：
- 仅实现 Piper 合成，去除 CosyVoice/Chatterbox。
- 兼容多种 Piper CLI 习惯：`piper` 或 `python3 -m piper`。
- 文本通过 stdin 传入，避免复杂转义。
- 自动回退多种命令模板；也支持在 YAML 里传 cmd_template 覆盖。
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import List, Dict, Optional, Tuple
import json
import os
import re
import shlex
import subprocess
import tempfile


@dataclass
class Speaker:
    name: str
    backend: str  # 固定为 "piper"
    model_path: Optional[str] = None
    config_path: Optional[str] = None
    sample_rate: int = 22050
    sentence_silence: float = 0.5
    tts_params: Dict = None
    cmd_template: Optional[str] = None  # 可选：自定义命令模板（高优先级）


class TTSClient:
    def __init__(self, speakers: List[Speaker], out_dir: str):
        self.speakers = speakers
        self.out_dir = out_dir
        os.makedirs(out_dir, exist_ok=True)

    # -------- 脚本切分：以「[角色] 开场...」风格为准 --------
    _line_re = re.compile(r"^\s*\[(?P<role>[^]]+)\]\s*(?P<text>.+?)\s*$")

    def _normalize_line(self, raw: str) -> Optional[Tuple[str, str]]:
        raw = raw.strip()
        if not raw or raw.startswith("#"):
            return None
        m = self._line_re.match(raw)
        if not m:
            return None
        role = m.group("role").strip()
        text = m.group("text").strip()
        return role, text

    def _pick_speaker(self, role_name: str) -> Speaker:
        for spk in self.speakers:
            if spk.name == role_name:
                return spk
        # 没有匹配就用第一个
        return self.speakers[0]

    def synthesize_dialogue(self, script: str, out_dir: str) -> List[str]:
        os.makedirs(out_dir, exist_ok=True)
        segs: List[str] = []
        idx = 0
        for raw in script.splitlines():
            parsed = self._normalize_line(raw)
            if not parsed:
                continue
            role, text = parsed
            spk = self._pick_speaker(role)
            idx += 1
            out_wav = os.path.join(out_dir, f"{idx:04d}_{role}.wav")
            try:
                self._piper_say(text, spk, out_wav)
            except Exception as e:
                print("❌ 合成失败 [piper] 角色:", role)
                print("文本:", text)
                print("错误:", repr(e))
                raise
            segs.append(out_wav)
        if not segs:
            raise RuntimeError("没有从脚本中解析到任何可合成的发言行。行格式示例：[主持人] 开场白……")
        return segs

    # -------------------- Piper --------------------
    def _piper_say(self, text: str, spk: Speaker, out_wav: str):
        if not spk.model_path or not os.path.exists(spk.model_path):
            raise FileNotFoundError(f"Piper 模型不存在：{spk.model_path}")
        # config.json 可选
        cfg = spk.config_path if (spk.config_path and os.path.exists(spk.config_path)) else None

        # 参数
        tts = spk.tts_params or {}
        ls = tts.get("length_scale")
        ns = tts.get("noise_scale")
        nw = tts.get("noise_w")
        sil = spk.sentence_silence

        # 构造参数片段（仅在不为 None 时拼接）
        def maybe(flag: str, val):
            if val is None:
                return ""
            return f" {flag} {shlex.quote(str(val))}"

        # 允许 YAML 自定义覆盖命令模板（高级用法）
        if spk.cmd_template:
            cmd = spk.cmd_template.format(
                model=spk.model_path,
                config=(cfg or ""),
                out=out_wav,
                text=text,
                length_scale=(ls if ls is not None else ""),
                noise_scale=(ns if ns is not None else ""),
                noise_w=(nw if nw is not None else ""),
                sentence_silence=sil,
                sr=spk.sample_rate,
            )
            self._run_cli_with_stdin(cmd, text)
            if not os.path.exists(out_wav):
                raise RuntimeError("自定义 Piper 命令执行后未生成输出 wav。")
            return

        # 标准模板（按顺序尝试）
        # 模板 A：piper 可执行文件在 PATH 内
        tmpl_A = (
            "piper "
            f"-m {shlex.quote(spk.model_path)}"
            + (f" -c {shlex.quote(cfg)}" if cfg else "")
            + f" -f {shlex.quote(out_wav)}"
            + maybe("--length_scale", ls)
            + maybe("--noise_scale", ns)
            + maybe("--noise_w", nw)
            + maybe("--sentence_silence", sil)
        )
        # 模板 B：python -m piper
        tmpl_B = tmpl_A.replace("piper ", "python3 -m piper ")

        # 模板 C：有些打包使用 --model/--config/--output_file 名称
        tmpl_C = (
            "piper "
            f"--model {shlex.quote(spk.model_path)}"
            + (f" --config {shlex.quote(cfg)}" if cfg else "")
            + f" --output_file {shlex.quote(out_wav)}"
            + maybe("--length_scale", ls)
            + maybe("--noise_scale", ns)
            + maybe("--noise_w", nw)
            + maybe("--sentence_silence", sil)
        )
        tmpl_D = tmpl_C.replace("piper ", "python3 -m piper ")

        last_err = None
        for cmd in (tmpl_A, tmpl_B, tmpl_C, tmpl_D):
            try:
                self._run_cli_with_stdin(cmd, text)
                if os.path.exists(out_wav):
                    return
                last_err = RuntimeError(f"命令成功返回但未生成文件：{cmd}")
            except Exception as e:
                last_err = e
                continue
        raise RuntimeError(f"Piper CLI 所有模板均失败。最后一个错误：{repr(last_err)}")

    def _run_cli_with_stdin(self, cmd: str, text: str):
        """
        通过 stdin 传入文本，避免复杂转义。
        """
        # 用 bash -lc 让管道/重定向语义一致
        bash_cmd = f"bash -lc {shlex.quote(cmd)}"
        # 将文本写入 stdin
        p = subprocess.Popen(
            bash_cmd,
            shell=True,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = p.communicate(input=text.encode("utf-8"))
        if p.returncode != 0:
            raise subprocess.CalledProcessError(p.returncode, bash_cmd, stdout, stderr)
